#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
使用MemoryClient的简化应用主模块
"""

import asyncio
import json
from pathlib import Path
from typing import Dict, List

from core.llm.factory import ModelType, create_llm_provider
from core.managers.dialogue_manager import DialogueManager
from core.utils.log import mylogger
from core.utils.mem0 import MemoryClient


class Mem0App:
    """使用MemoryClient的应用主类"""

    def __init__(self):
        self.input_file_path = Path(__file__).parent.parent / "data" / "input.json"
        self.model_types = [ModelType.DOUBAO_FLASH]
        self.model_names = {
            ModelType.DOUBAO_FLASH: "doubaoFlash",
        }

    def _load_input_data(self) -> List[Dict]:
        """从input.json加载输入数据"""
        try:
            with open(self.input_file_path, "r", encoding="utf-8") as f:
                input_data = json.load(f)
            mylogger.info(f"成功加载输入文件: {self.input_file_path}")
            return input_data if isinstance(input_data, list) else [input_data]
        except (FileNotFoundError, json.JSONDecodeError) as e:
            mylogger.error(f"无法读取输入文件 {self.input_file_path}: {e}")
            raise

    def _save_result_data(self, result_data: List[Dict], model_type: ModelType):
        """保存结果数据到特定模型的output.json"""
        try:
            model_name = self.model_names[model_type]
            output_path = Path(__file__).parent.parent / "data" / f"{model_name}_output.json"
            output_path.parent.mkdir(parents=True, exist_ok=True)

            with open(output_path, "w", encoding="utf-8") as f:
                json.dump(result_data, f, ensure_ascii=False, indent=4)
            mylogger.info(f"成功保存结果到: {output_path}")
        except Exception as e:
            mylogger.error(f"保存结果文件失败: {e}")
            raise

    async def _get_relevant_memories(self, query: str, agent_id: str, limit: int = 3) -> str:
            """获取相关记忆并格式化为字符串"""
            try:
                # 获取搜索结果
                related_memories = MemoryClient.search(
                    query=query,
                    agent_id=agent_id,
                    limit=limit
                )
                if isinstance(related_memories, str):
                    related_memories = json.loads(related_memories)

                mylogger.info(f"相关记忆: {related_memories}")

                results = related_memories.get("results", [])
                if not results:
                    return "暂无相关记忆"

                memory_texts = []
                for memory in results:
                    memory_content = memory.get('memory', '')
                    metadata = memory.get('metadata', {})
                    if memory_content:
                        memory_texts.append(f"- {memory_content} (metadata: {metadata})")
                
                mylogger.info(f"memory_texts: {memory_texts}")
                result = "\n".join(memory_texts) if memory_texts else "暂无相关记忆"
                return result

            except Exception as e:
                mylogger.error(f"获取记忆失败: {e}")
                return "暂无相关记忆"

    async def _add_memory(self, content: str, agent_id: str, role: str = "user", memory_type: str = "conversation") -> bool:
        """添加记忆到MemoryClient"""
        try:
            messages = [{"role": role, "content": content}]
            mylogger.info(f"添加记忆: {messages}")
            result = MemoryClient.add(
                messages=messages,
                agent_id=agent_id,
                metadata={"type": memory_type}
            )
            
            mylogger.info(f"添加记忆结果: {result}")
            return True
        except Exception as e:
            mylogger.error(f"添加记忆失败: {e}")
            return False

    async def process_dialogue_with_memory(self, npc_name: str):
        """处理多轮对话，使用MemoryClient进行记忆管理"""
        input_data_list = self._load_input_data()
        
        for model_type in self.model_types:
            model_name = self.model_names[model_type]
            mylogger.info(f"使用模型 {model_name} 处理对话")

            # 创建组件
            llm = create_llm_provider(model_type=model_type)
            dialogue_manager = DialogueManager(npc_name, llm, model_name)

            result_data = []

            for input_item in input_data_list:
                # 判断input和skill
                user_input = input_item.get("input", "")
                user_skill = input_item.get("skill", "")
                
                if user_input:
                    content = user_input
                    memory_type = "conversation"
                elif user_skill:
                    content = user_skill  
                    memory_type = "event"
                else:
                    continue

                mylogger.info(f"处理{memory_type}: {content}")

                # 在对话前获取相关记忆并注入到系统提示中
                memory_context = await self._get_relevant_memories(content, npc_name)
                dialogue_manager.template_params["memory_context"] = memory_context

                # 生成响应
                response = await dialogue_manager.generate_response(content)
                print(f"NPC {npc_name} ({model_name}) 回复: {response}")

                # 在对话后添加记忆
                await self._add_memory(content, npc_name, role="user", memory_type=memory_type)

                result_data.append({"input": content, "response": response, "type": memory_type})

            self._save_result_data(result_data, model_type)

        return "处理完成"


async def main():
    """主函数"""
    app = Mem0App()
    npc_name = "赵凤儿"
    
    mylogger.info("开始处理多轮对话记忆交互")
    await app.process_dialogue_with_memory(npc_name)
    
    mylogger.info("处理完成")


if __name__ == "__main__":
    asyncio.run(main())